import numpy as np
import graphviz

# 定义树节点类
class TreeNode:
    def __init__(self, feature=None, threshold=None, left=None, right=None, *, value=None):
        self.feature = feature  # 分割特征（对于一维数据，特征索引为0）
        self.threshold = threshold  # 分割阈值
        self.left = left  # 左子树                         
        self.right = right  # 右子树
        self.value = value  # 叶节点的预测值

# 计算MSE
def mse_unit(y_true, y_pred):
    return (y_true - y_pred) ** 2
def mse(y_list):
    y_avg = np.mean(y_list)
    mse_val = 0
    for y in y_list:
        mse_val += mse_unit(y, y_avg)
    return mse_val

# 寻找最佳分割点
def find_best_split(X, y):
    best_loss = float('inf')
    best_threshold = None

    # 获取所有唯一的x值作为候选分割点
    thresholds = np.unique(X)

    for threshold in thresholds:
        # 左子集：x <= threshold
        left_idx = X <= threshold
        # 右子集：x > threshold
        right_idx = X > threshold

        loss_left = mse(y[left_idx])
        loss_right = mse(y[right_idx])
        loss = loss_right + loss_left
        if loss < best_loss:
            best_loss = loss
            best_threshold = threshold

    return best_loss, best_threshold

# 递归构建回归树
def build_tree(X, y, min_samples_split=2, max_depth=None, depth=0):
    # 如果样本数为1，创建叶节点
    if X.size == 1:
        leaf_value = np.mean(y)
        return TreeNode(value=leaf_value)

    # 寻找最佳分割点
    loss, threshold = find_best_split(X, y)

    # 如果无法分割或分割后子集过小，创建叶节点
    if loss == -1:
        leaf_value = np.mean(y)
        return TreeNode(value=leaf_value)

    # 分割数据集
    left_idx = X <= threshold
    right_idx = X > threshold

    # 递归构建左右子树
    left = build_tree(X[left_idx], y[left_idx], min_samples_split, max_depth, depth + 1)
    right = build_tree(X[right_idx], y[right_idx], min_samples_split, max_depth, depth + 1)

    # 创建内部节点
    return TreeNode(feature=0, threshold=threshold, left=left, right=right)

def print_tree(node, depth=0):
    if node is None:
        return

    indent = "  " * depth
    if node.value is not None:
        print(f"{indent}Leaf: Value = {node.value}")
    else:
        print(f"{indent}Node: Feature = {node.feature}, Threshold = {node.threshold}")
        print(f"{indent}Left:")
        print_tree(node.left, depth + 1)
        print(f"{indent}Right:")
        print_tree(node.right, depth + 1)

# 绘制回归树
def plot_tree(node, dot=None):
    if dot is None:
        dot = graphviz.Digraph()
        dot.node(name=str(id(node)), label=f"Threshold: {node.threshold}\nValue: {node.value}")

    if node.left:
        dot.node(name=str(id(node.left)), label=f"Threshold: {node.left.threshold}\nValue: {node.left.value}")
        dot.edge(str(id(node)), str(id(node.left)), label=f"x <= {node.threshold}")
        plot_tree(node.left, dot=dot)

    if node.right:
        dot.node(name=str(id(node.right)), label=f"Threshold: {node.right.threshold}\nValue: {node.right.value}")
        dot.edge(str(id(node)), str(id(node.right)), label=f"x > {node.threshold}")
        plot_tree(node.right, dot=dot)

    return dot

# 题目数据
X = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.array([4.50, 4.75, 4.91, 5.34, 5.80, 7.05, 7.90, 8.23, 8.70, 9.00])

# 构建回归树
root = build_tree(X, y, min_samples_split=2, max_depth=None)

# 绘制回归树并保存为PNG文件
dot = plot_tree(root)
dot.render('regression_tree', format='png', cleanup=True)
print_tree(root)